import torch
import random

def synthetic_data(w, b, num_examples): #@save
    """生成y=Xw+b+噪声"""
    # 生成特征矩阵 X： num_examples 行， len(w) 列， 元素服从 N(0,1)
    # torch.normal(均值, 标准差, 张量形状)
    # 当前均值为标量0，即生成的所有元素均值为0
    # 当前标准差为标量1，即生成的所有元素标准差为1
    X = torch.normal(0, 1, (num_examples, len(w)))
    # 计算理论标签： y = X@w + b (矩阵乘法 + 偏置值)
    # torch.matmul() 实现矩阵乘法
    y = torch.matmul(X, w) + b
    # 添加噪声
    y += torch.normal(0, 0.01, y.shape)
    # 重塑标签为列向量，返回特征 X、标签 y
    return X, y.reshape((-1, 1))

def data_iter(batch_size, features, labels):
    num_examples = len(features)
    indices = list(range(num_examples))
    # 这些样本是随机读取的，没有特定的顺序
    random.shuffle(indices)
    for i in range(0, num_examples, batch_size):
        batch_indices = torch.tensor(indices[i: min(i + batch_size, num_examples)])
        # yield 语句可以让一个函数变成一个生成器
        yield features[batch_indices], labels[batch_indices]

if __name__ == "__main__":
    from matplotlib import pyplot as plt
    true_w = torch.tensor([2, -3.4])
    true_b = 4.2
    features, labels = synthetic_data(true_w, true_b, 1000)

    print('features:', features[0],'\nlabel:', labels[0])

    plt.figure(figsize=(6, 4))
    # detach() 返回一个新的张量，从当前计算图中分离出来，requires_grad=False
    # numpy() 将张量转换为 NumPy 数组
    plt.scatter(features[:, (1)].detach().numpy(), labels.detach().numpy(), 1)
    plt.show()

    counter = 0
    for X, y in data_iter(10, features, labels):
        print(X, '\n', y)
        counter += 1
        if counter == 3:
            break
    input('Press Enter to continue...')  # 防止图形一闪而过